package com.central.gateway.filter;

import com.alibaba.fastjson.JSONObject;
import com.central.gateway.common.AuthorizationBo;
import com.central.gateway.common.Constants;
import com.central.gateway.common.LogVo;
import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.HttpHeaders;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.Base64;
import java.util.Objects;

import static java.nio.charset.StandardCharsets.UTF_8;

/**
 * @author zouzhihui
 */
@Component
@Slf4j
public class AddAuthorizationRequestHeaderGatewayFilterFactory extends AbstractGatewayFilterFactory<AbstractGatewayFilterFactory.NameConfig> {

    public AddAuthorizationRequestHeaderGatewayFilterFactory() {
        super(AbstractGatewayFilterFactory.NameConfig.class);
    }

    @Override
    public GatewayFilter apply(NameConfig config) {
        return new GatewayFilter() {
            @Override
            public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
                LogVo logVo = (LogVo) exchange.getAttributes().get(Constants.LOG_VO_ATTR);
                if (Objects.nonNull(logVo)) {
                    AuthorizationBo authorizationBo = JSONObject.parseObject(logVo.getBodyParameter(), AuthorizationBo.class);
                    byte[] encode = Base64.getEncoder().encode(String.format("%s:%s", authorizationBo.getMerchant(),
                            authorizationBo.getSecret()).getBytes(UTF_8));
                    String authorization = new String(encode, UTF_8);
                    ServerHttpRequest request = exchange.getRequest().mutate()
                            .header("Authorization", "Basic " + authorization).build();
                    return chain.filter(exchange.mutate().request(request).build());
                } else {
                    return DataBufferUtils.join(exchange.getRequest().getBody()).flatMap(dataBuffer -> {
                        byte[] bytes = new byte[dataBuffer.readableByteCount()];
                        dataBuffer.read(bytes);
                        String bodyString = new String(bytes, UTF_8);
                        DataBufferUtils.release(dataBuffer);
                        Flux<DataBuffer> cachedFlux = Flux.defer(() -> {
                            DataBuffer buffer = exchange.getResponse().bufferFactory().wrap(bytes);
                            return Mono.just(buffer);
                        });

                        AuthorizationBo authorizationBo = JSONObject.parseObject(bodyString, AuthorizationBo.class);
                        byte[] encode = Base64.getEncoder().encode(String.format("%s:%s", authorizationBo.getMerchant(),
                                authorizationBo.getSecret()).getBytes(UTF_8));
                        String authorization = new String(encode, UTF_8);

                        ServerHttpRequest mutatedRequest = new ServerHttpRequestDecorator(exchange.getRequest()) {
                            @Override
                            public Flux<DataBuffer> getBody() {
                                return cachedFlux;
                            }

                            @Override
                            public HttpHeaders getHeaders() {
                                HttpHeaders headers = new HttpHeaders();
                                headers.addAll(super.getHeaders());
                                headers.add("Authorization", "Basic " + authorization);
                                return headers;
                            }
                        };


                        return chain.filter(exchange.mutate().request(mutatedRequest).build());
                    });
                }
            }
        };
    }
}
